#include <iostream>
#include <string>
#include <cstring>
#include <cstdlib>
#include <ctime>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <sys/time.h>
#include <unistd.h>

#include "MRKeys.h"

NTL_CLIENT;

MRKeys::MRKeys(MRProcessArgs &mrp) {
	keySize = 4096;
	b64_chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
}

MRKeys::~MRKeys() {

}

unsigned long long MRKeys::getSeedNum(void) {
	struct timeval now;

	if(gettimeofday(&now, NULL) < 0) {
		perror("gettimeofday");
		exit(EXIT_FAILURE);
	}

	unsigned long long seedNum = (unsigned long long) 
		(now.tv_sec * 1000 * 1000) +
		(now.tv_usec);

	return seedNum;
}

void MRKeys::byte64Encode(unsigned char *asc, 
		       unsigned char *bin,
		       unsigned len) {
	const int rem = len % 3;
	unsigned alen;

	if((len * 4) % 3 == 0)
		alen  = len * 4 / 3;
	else
		alen = len * 4 / 3 + 1;

	unsigned char temp[alen];
	memset(temp, 0, alen);
	int last_i, last_j;

	for (int i = 0, j = 0; i < len; i += 3, j += 4) {
		temp[j] = (bin[i] & 0xfc) >> 2;
		temp[j + 1] = ((bin[i] & 0x3) << 4) + ((bin[i + 1] & 0xf0) >> 4);
		temp[j + 2] = ((bin[i + 1] & 0xf) << 2) + ((bin[i + 2] & 0xc0) >> 6);
		temp[j + 3] = (bin[i + 2] & 0x3f);

		last_i = i;
		last_j = j;
	}

	if (rem == 2) {
		temp[last_j]     = (bin[last_i] & 0xfc) >> 2;
		temp[last_j + 1] = ((bin[last_i] & 0x3) << 4) + 
			((bin[last_i + 1] & 0xf0) >> 4);
		temp[last_j + 2] = (bin[last_i + 1] & 0xf);
	}
	else if (rem == 1) {
		temp[last_j] = (bin[last_i] & 0xfc) >> 2;
		temp[last_j + 1] = (bin[last_i] & 0x3);
	}


#ifdef MR_DEBUG
	cout << "(byte64Encode)"
	     << "\trem: " << rem 
	     << "\talen: "   << alen 
	     << "\tlast_j: " << last_j 
	     << "\tlen: "   << len 
	     << "\tlast_i: "  << last_i 
	     << endl;
#endif // MR_DEBUG

	int i;
	for (i = 0; i < alen; i++) { 
		asc[i] = b64_chars[(unsigned)temp[i]];
	}
}

int MRKeys::byte64Decode(unsigned char *bin, 
		       unsigned char *asc,
		       unsigned alen) {
	unsigned blen = (alen * 3) / 4;            // Raw byte string
						   // length
	const unsigned rem = alen % 4;             // 4 ASCII
						   // characters eq. 3 bytes

	int last_i, last_j;
	memset(bin, 0, keySize);

	for(int i = 0, j = 0; i < blen; i += 3, j += 4) {
		unsigned char c1 = b64_chars.find((unsigned) asc[j]);
		unsigned char c2 = b64_chars.find((unsigned) asc[j + 1]);
		unsigned char c3 = b64_chars.find((unsigned) asc[j + 2]);
		unsigned char c4 = b64_chars.find((unsigned) asc[j + 3]);

		bin[i] = (c1 << 2) + ((c2 & 0x30) >> 4);
		bin[i + 1] = ((c2 & 0xf) << 4) + ((c3 & 0x3c) >> 2);
		bin[i + 2] = ((c3 & 0x3) << 6) + (c4 & 0x3f);

		last_i = i;
		last_j = j;
	}

	if(rem != 0) {
		unsigned char c1 = b64_chars.find((unsigned) asc[last_j]);
		unsigned char c2 = b64_chars.find((unsigned) asc[last_j + 1]);

		if(rem == 3) {
			unsigned char c3 = b64_chars.find((unsigned) 
							  asc[last_j + 2]);

			bin[last_i] = (c1 << 2) + ((c2 & 0x30) >> 4);
			bin[last_i + 1] = ((c2 & 0xf) << 4) + (c3  & 0xf);
 			bin[last_i + 2] = '\0';
		}
		else if (rem == 2) { 
			bin[last_i] = (c1 << 2) + (c2 & 0x3);
		}
	}

#ifdef MR_DEBUG
	cout << "(byte64Decode)"
	     << "\trem: " << rem 
	     << "\talen: "   << alen 
	     << "\tlast_j: " << last_j 
	     << "\tblen: "   << blen 
	     << "\tlast_i: "  << last_i 
	     << endl;
#endif // MR_DEBUG

	return blen;
}

int MRKeys::genKeys() {
	ZZ seed;
	seed  = (unsigned long long)this->getSeedNum();

	SetSeed(seed);

// 	p = RandomLen_ZZ(keySize);
// 	q = RandomLen_ZZ(keySize);

 	p = GenPrime_ZZ(keySize);
 	q = GenPrime_ZZ(keySize);

	n = p * q;
	totient_n = (p - 1) * (q - 1);

	ZZ e, d;
	e = d = (long) 0;

	while((e < 1) || (e > totient_n) || (GCD(e, totient_n) != 1))  {
		e  = RandomLen_ZZ(keySize);
	}

	if (InvModStatus(d, e, totient_n) != 0) {
		cerr << "Error: InvModStatus!" << endl;
		exit(1);
	}

	unsigned char publicKeyBytes[keySize], privateKeyBytes[keySize];
	const unsigned pubKeyByteLen = NumBytes(e);
	const unsigned priKeyByteLen = NumBytes(d);

	BytesFromZZ(publicKeyBytes,  e, pubKeyByteLen);
	BytesFromZZ(privateKeyBytes, d, priKeyByteLen);

	unsigned char armoredPubText[keySize];
	unsigned char armoredPriText[keySize];

	byte64Encode(armoredPubText, publicKeyBytes,  pubKeyByteLen);
 	byte64Encode(armoredPriText, privateKeyBytes, priKeyByteLen);

	unsigned char pubBytes[keySize], priBytes[keySize];;
	int bl_pub = byte64Decode(pubBytes, 
				armoredPubText, 
				strlen((char *)armoredPubText));

	int bl_pri = byte64Decode(priBytes, 
				armoredPriText, 
				strlen((char *)armoredPriText));

	ZZ a, b;
	ZZFromBytes(a, pubBytes, bl_pub);
	ZZFromBytes(b, priBytes, bl_pri);

	int cmp = memcmp((char *)publicKeyBytes, (char *)pubBytes, 512);
	if( cmp != 0) {
		cerr << "Error: Byte String mismatch!" << endl;
		cout << "Original: " << publicKeyBytes << endl;
		cout << "Modified: " << pubBytes << endl;
		cout << "CMP: " << cmp << endl;
	}

	ZZ diff_pub = e - a;
	ZZ diff_pri = d - b;
	ZZ diff = diff_pub + diff_pri;

	if (diff_pub != 0)
		cout << "Public Not Same: " << diff_pub << endl;

	if (diff_pri != 0)
		cout << "Private Not Same: " << diff_pri << endl;

	int priFileDesc = open("./key.pri", O_CREAT | O_WRONLY | O_SYNC, 
			       S_IWUSR | S_IRUSR);
	int pubFileDesc = open("./key.pub", O_CREAT | O_WRONLY | O_SYNC, 
			       S_IWUSR | S_IRUSR);

	if ((priFileDesc < 0) || (pubFileDesc < 0)) {
		perror("open");
		exit(EXIT_FAILURE);
	}

	if (write(pubFileDesc, armoredPubText, strlen((char *)armoredPubText)) < 0) {
		perror("write");
		exit(EXIT_FAILURE);
	}

	if (write(priFileDesc, armoredPriText, strlen((char *)armoredPriText)) < 0) {
		perror("write");
		exit(EXIT_FAILURE);
	}

	close(pubFileDesc);
	close(priFileDesc);
}

